
import torch.nn as nn
from utils import Params
from models.language_model import *
from models.encoder_model import *
from models.classify_model import *
class NERTokenClassification(nn.Module):
    def __init__(self, params: Params):
        super(NERTokenClassification, self).__init__()

        # 创建三个具体层的模型
        self.emb_layer = eval(params.lm_layer_name)(params)
        self.encoder_layer = eval(params.encoder_layer_name)(params)
        self.classify_layer = eval(params.classify_layer_name)(params)

    def forward(self, input_ids, input_masks, labels=None):
        z = self.emb_layer(input_ids, input_masks)
        z = self.encoder_layer(z, input_masks)
        z = self.classify_layer(z, input_masks, labels=labels)
        return z

